import os
import pickle

def get_root_dir(path):
    path_list=path.split(os.path.sep)
    index=path_list.index("featurelib")
    return os.path.sep.join(path_list[:index+1])

class SmsEnLoan0V1NerModel(object):

    def __init__(self):
        self.ROOT_DIR = get_root_dir(os.path.abspath("."))
        self.CONF_DIR = os.path.join(
            self.ROOT_DIR, "feature_conf", "sms", "un", "sms_en_loan0_v1"
        )
        self.MODEL_CONF_DIR = os.path.join(self.CONF_DIR, "model_conf")
        pkl_name = "english_cmp_ner.pkl"
        with open(os.path.join(self.MODEL_CONF_DIR, pkl_name), "rb") as f:
            self.nlp = pickle.load(f)

    def predict(self, msg_list):
        doc_list = self.nlp.pipe(msg_list)
        res = []
        for doc in doc_list:
            entities = {}
            for ent in doc.ents:
                entities[ent.label_] = ent.text
            res.append(entities)
        return res